package rs.fon.neuroph.regression;

import java.util.ArrayList;
import java.util.Iterator;

import org.neuroph.core.Connection;
import org.neuroph.core.NeuralNetwork;
import org.neuroph.core.Neuron;
import org.neuroph.core.learning.IterativeLearning;
import org.neuroph.core.learning.SupervisedTrainingElement;
import org.neuroph.core.learning.TrainingElement;
import org.neuroph.core.learning.TrainingSet;
import org.neuroph.core.transfer.Linear;
import org.neuroph.nnet.comp.InputNeuron;

import com.rapidminer.example.Attribute;
import com.rapidminer.example.Example;
import com.rapidminer.example.ExampleSet;
import com.rapidminer.operator.Model;
import com.rapidminer.operator.OperatorException;

public class NeurophRegressionAdapter {

	private int maxIterations;
	
	public NeurophRegressionAdapter(int maxIterations) {
		this.maxIterations = maxIterations;
	}
	
	public Model trainNNModel(String filePath, ExampleSet exampleSet) throws OperatorException {
		//prepare
		TrainingSet<TrainingElement> ts = convertExampleSetToTrainingSet(exampleSet);
		NeuralNetwork nnet = NeuralNetwork.load(filePath);

		//Check if the NN is compatible with the given dataset
		checkNNCompatibility(exampleSet, nnet);

		//learn
		IterativeLearning learning = (IterativeLearning)nnet.getLearningRule();
		learning.setMaxIterations(maxIterations);
		nnet.learn(ts);

		//return model
		Model resultModel = new NeurophRegressionModel(exampleSet,nnet);
		return resultModel;
	}



	public static TrainingSet<TrainingElement> convertExampleSetToTrainingSet(ExampleSet es) throws OperatorException{

		if (es.getAttributes().getLabel()==null)
			throw new OperatorException("No label defined in the training set");

		TrainingSet<TrainingElement> ts = new TrainingSet<TrainingElement>();
		for (Example e: es){
			ArrayList<Double> values = new ArrayList<Double>();
			for (Iterator<Attribute> i=es.getAttributes().iterator(); i.hasNext(); )
				values.add(e.getNumericalValue(i.next()));
			ArrayList<Double> labels = new ArrayList<Double>();
			labels.add(e.getNumericalValue(e.getAttributes().getLabel()));
			ts.addElement(new SupervisedTrainingElement(values, labels));
		}

		return ts;
	}


	private void checkNNCompatibility(ExampleSet exampleSet, NeuralNetwork nnet) throws OperatorException {

		int numberOfInputAttributes=0; 
		for (Iterator<Attribute> i=exampleSet.getAttributes().iterator(); i.hasNext(); ){
			i.next();
			numberOfInputAttributes++;
		}

		int inputNeuronsMissing = numberOfInputAttributes - nnet.getInputNeurons().size();
		if (inputNeuronsMissing>0){
			Neuron sampleNeuron = nnet.getInputNeurons().get(0);
			for (int i=0; i<inputNeuronsMissing; i++){
				InputNeuron n = new InputNeuron();
				for (Connection c: sampleNeuron.getOutConnections()){
					c.getToNeuron().addInputConnection(n);
				}
				sampleNeuron.getParentLayer().addNeuron(n);
				nnet.getInputNeurons().add(n);
			}
		}

		if (inputNeuronsMissing<0)
			for (int i=0; i<-inputNeuronsMissing; i++)
				nnet.getInputNeurons().remove(0);

		nnet.getOutputNeurons().get(0).setTransferFunction(new Linear());

		if (numberOfInputAttributes != nnet.getInputNeurons().size())
			throw new OperatorException("Number of input neurons does not match the number of attributes");
		if (nnet.getOutputNeurons().size() != 1)
			throw new OperatorException("Number of output neurons should be 1 for the regression task");
	}


}
